{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "c75b2353-5dce-47d1-b841-ac36b918d419",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-10-11 07:12:09.388515: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.\n",
      "To enable the following instructions: AVX2 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.\n",
      "2023-10-11 07:12:10.483964: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2023-10-11 07:12:12,827] [INFO] [real_accelerator.py:158:get_accelerator] Setting ds_accelerator to cuda (auto detect)\n"
     ]
    }
   ],
   "source": [
    "import json\n",
    "import random\n",
    "import numpy as np\n",
    "import torch\n",
    "from datasets import load_dataset, Dataset\n",
    "from tqdm import tqdm\n",
    "from transformers import AutoModel, AutoTokenizer, pipeline\n",
    "\n",
    "\n",
    "from repe import repe_pipeline_registry\n",
    "repe_pipeline_registry()\n",
    "\n",
    "# Configurationsrandom.seed(0)\n",
    "np.random.seed(0)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "9ca4b8af-57d4-486d-8966-b7bb8f97e576",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n",
      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n"
     ]
    }
   ],
   "source": [
    "model_name_or_path =  \"microsoft/deberta-xxlarge-v2-mnli\"\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)\n",
    "model = AutoModel.from_pretrained(model_name_or_path).half().cuda()\n",
    "\n",
    "rep_pipeline =  pipeline(\"rep-reading\", model=model, tokenizer=tokenizer, device=model.device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "3b636f9d-7b55-4009-868a-3bf0fa521f7b",
   "metadata": {},
   "outputs": [],
   "source": [
    "n_difference = 1\n",
    "hidden_layers = list(range(-1, -model.config.num_hidden_layers, -1))\n",
    "direction_method = \"pca\"\n",
    "direction_finder_kwargs= {\"n_components\": 1}\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "04645393-d6f3-44cb-90ea-2629e2e5c127",
   "metadata": {},
   "source": [
    "## RTE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "3cd2de19-b66a-4080-94d5-723112601708",
   "metadata": {},
   "outputs": [],
   "source": [
    "template_str = \"Consider the {concept} of the sentences: Hypothesis: {sent1}\\Premise: {sent2}\"\n",
    "rep_token = 3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "17c4b04e-5bbb-4e26-93c7-9ebcbce7e7dd",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Found cached dataset glue (/data/long_phan/.cache/huggingface/datasets/glue/rte/1.0.0/dacbe3125aa31d7f70367a07a8a9e72a5a0bfeb5fc42e75c9db75b96da6053ad)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d3dd2eabe7274c288ed330efff079c7c",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached shuffled indices for dataset at /data/long_phan/.cache/huggingface/datasets/glue/rte/1.0.0/dacbe3125aa31d7f70367a07a8a9e72a5a0bfeb5fc42e75c9db75b96da6053ad/cache-dddfb666e65c61ad.arrow\n"
     ]
    }
   ],
   "source": [
    "def samples(ds, test_set=False):\n",
    "    pairs = []\n",
    "    for e in ds:\n",
    "        concepts = ['entailment', 'contradiction']\n",
    "                \n",
    "        if e['label'] == 1 and test_set: # flip the concepts in test for easy eval\n",
    "            concepts = concepts[::-1]\n",
    "    \n",
    "        pair = [template_str.format(concept=c, sent1=e['sentence1'], sent2=e['sentence2']) for c in concepts]\n",
    "        pairs.extend(pair)\n",
    "    return pairs\n",
    "\n",
    "dataset = load_dataset(\"glue\", \"rte\")\n",
    "max_train_samples = 64\n",
    "train_dataset, test_dataset = dataset['train'], dataset['validation']\n",
    "train_dataset =  train_dataset.shuffle(seed=0)\n",
    "train_dataset = train_dataset.select(range(max_train_samples))\n",
    "\n",
    "train_data, train_labels = samples(train_dataset), train_dataset['label']\n",
    "train_labels = [[1, 0] if label == 0 else [0, 1] for label in train_labels]\n",
    "test_data =  samples(test_dataset, test_set=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "ca99bd80-b9d4-4651-be99-6e623428b12c",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-1: 0.8989169675090253\n",
      "=====\n",
      "-2: 0.6209386281588448\n",
      "=====\n",
      "-3: 0.7725631768953068\n",
      "=====\n",
      "-4: 0.8953068592057761\n",
      "=====\n",
      "-5: 0.8700361010830325\n",
      "=====\n",
      "-6: 0.7581227436823105\n",
      "=====\n",
      "-7: 0.7292418772563177\n",
      "=====\n",
      "-8: 0.6931407942238267\n",
      "=====\n",
      "-9: 0.628158844765343\n",
      "=====\n",
      "-10: 0.7364620938628159\n",
      "=====\n",
      "-11: 0.7220216606498195\n",
      "=====\n",
      "-12: 0.7509025270758123\n",
      "=====\n",
      "-13: 0.6967509025270758\n",
      "=====\n",
      "-14: 0.7184115523465704\n",
      "=====\n",
      "-15: 0.6823104693140795\n",
      "=====\n",
      "-16: 0.7509025270758123\n",
      "=====\n",
      "-17: 0.5631768953068592\n",
      "=====\n",
      "-18: 0.516245487364621\n",
      "=====\n",
      "-19: 0.5018050541516246\n",
      "=====\n",
      "-20: 0.48375451263537905\n",
      "=====\n",
      "-21: 0.5126353790613718\n",
      "=====\n",
      "-22: 0.4981949458483754\n",
      "=====\n",
      "-23: 0.51985559566787\n",
      "=====\n",
      "-24: 0.5054151624548736\n",
      "=====\n",
      "-25: 0.47653429602888087\n",
      "=====\n",
      "-26: 0.4548736462093863\n",
      "=====\n",
      "-27: 0.5054151624548736\n",
      "=====\n",
      "-28: 0.516245487364621\n",
      "=====\n",
      "-29: 0.4729241877256318\n",
      "=====\n",
      "-30: 0.4368231046931408\n",
      "=====\n",
      "-31: 0.4584837545126354\n",
      "=====\n",
      "-32: 0.4404332129963899\n",
      "=====\n",
      "-33: 0.44404332129963897\n",
      "=====\n",
      "-34: 0.4368231046931408\n",
      "=====\n",
      "-35: 0.44404332129963897\n",
      "=====\n",
      "-36: 0.4584837545126354\n",
      "=====\n",
      "-37: 0.4404332129963899\n",
      "=====\n",
      "-38: 0.4368231046931408\n",
      "=====\n",
      "-39: 0.4404332129963899\n",
      "=====\n",
      "-40: 0.4332129963898917\n",
      "=====\n",
      "-41: 0.4368231046931408\n",
      "=====\n",
      "-42: 0.5379061371841155\n",
      "=====\n",
      "-43: 0.51985559566787\n",
      "=====\n",
      "-44: 0.4548736462093863\n",
      "=====\n",
      "-45: 0.51985559566787\n",
      "=====\n",
      "-46: 0.5270758122743683\n",
      "=====\n",
      "-47: 0.4693140794223827\n",
      "=====\n"
     ]
    }
   ],
   "source": [
    "batch_size=128\n",
    "rep_reader = rep_pipeline.get_directions(\n",
    "    train_data,\n",
    "    rep_token=rep_token, \n",
    "    hidden_layers=hidden_layers, \n",
    "    n_difference=n_difference, \n",
    "    train_labels=train_labels, \n",
    "    direction_method=direction_method,\n",
    "    direction_finder_kwargs=direction_finder_kwargs,\n",
    "    batch_size=batch_size,\n",
    "    \n",
    "    max_length=2048,\n",
    "    padding=\"longest\",\n",
    ")\n",
    "\n",
    "batch_size=128\n",
    "results_val = {layer: {} for layer in hidden_layers}\n",
    "\n",
    "H_tests = rep_pipeline(\n",
    "                    test_data,\n",
    "                    rep_token=rep_token, \n",
    "                    hidden_layers=hidden_layers, \n",
    "                    rep_reader=rep_reader,\n",
    "                    batch_size=batch_size,\n",
    "                    max_length=2048,\n",
    "                    padding=\"longest\")\n",
    "\n",
    "n_choices=2\n",
    "for layer in hidden_layers:\n",
    "    H_test = [H[layer] for H in H_tests] \n",
    "    H_test = [H_test[i:i+n_choices] for i in range(0, len(H_test), n_choices)]\n",
    "\n",
    "    sign = rep_reader.direction_signs[layer]\n",
    "    eval_func = min if sign == -1 else max\n",
    "    \n",
    "    cors = np.mean([eval_func(H) == H[0] for H in H_test])\n",
    "    results_val[layer] = cors\n",
    "    \n",
    "    print(f\"{layer}: {cors}\")\n",
    "    print(\"=====\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fc7a1bea-a6bb-483e-a344-cb22e5fb9f7c",
   "metadata": {},
   "source": [
    "## Boolq"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "85fc0ec1-a958-4acc-8f2a-5f2e857d5bc8",
   "metadata": {},
   "outputs": [],
   "source": [
    "template_str = \"Consider the {concept} of answering Yes to the question:\\nQuestion: {question}?\\nContext: {context}\"\n",
    "rep_token = 3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "eb7fc9db-139c-40ba-adea-7464ca3d3061",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Found cached dataset boolq (/data/long_phan/.cache/huggingface/datasets/boolq/default/0.1.0/bf0dd57da941c50de94ae3ce3cef7fea48c08f337a4b7aac484e9dddc5aa24e5)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "45470a46bec041ba9a5ab470d669ea78",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/2 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached shuffled indices for dataset at /data/long_phan/.cache/huggingface/datasets/boolq/default/0.1.0/bf0dd57da941c50de94ae3ce3cef7fea48c08f337a4b7aac484e9dddc5aa24e5/cache-deec76edf279b5b0.arrow\n"
     ]
    }
   ],
   "source": [
    "def samples(ds, test_set=False):\n",
    "    pairs = []\n",
    "    for e in ds:\n",
    "        concepts = ['correctness', 'incorrectness']\n",
    "        question = e['question']\n",
    "        context = e['passage']\n",
    "                \n",
    "        if e['answer'] == False and test_set: # flip the concepts in test for easy eval\n",
    "            concepts = concepts[::-1]\n",
    "\n",
    "        pair = [template_str.format(concept=c, question=question, context=context) for c in concepts]\n",
    "        pairs.extend(pair)\n",
    "    return pairs\n",
    "\n",
    "dataset = load_dataset(\"boolq\")\n",
    "max_train_samples = 64\n",
    "train_dataset, test_dataset = dataset['train'], dataset['validation']\n",
    "train_dataset =  train_dataset.shuffle(seed=0)\n",
    "train_dataset = train_dataset.select(range(max_train_samples))\n",
    "test_dataset = test_dataset.select(range(500)) # comment out for full set\n",
    "\n",
    "train_data = samples(train_dataset)\n",
    "train_labels = [[1,0] if l else [0,1] for l in train_dataset['answer']] # we get the index of true answer to decide the signs here, not for train\n",
    "test_data =  samples(test_dataset, test_set=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "22f90c8c-19bc-4970-af9e-96976d59b93f",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-1: 0.788\n",
      "=====\n",
      "-2: 0.626\n",
      "=====\n",
      "-3: 0.794\n",
      "=====\n",
      "-4: 0.788\n",
      "=====\n",
      "-5: 0.776\n",
      "=====\n",
      "-6: 0.766\n",
      "=====\n",
      "-7: 0.744\n",
      "=====\n",
      "-8: 0.748\n",
      "=====\n",
      "-9: 0.748\n",
      "=====\n",
      "-10: 0.752\n",
      "=====\n",
      "-11: 0.762\n",
      "=====\n",
      "-12: 0.768\n",
      "=====\n",
      "-13: 0.756\n",
      "=====\n",
      "-14: 0.772\n",
      "=====\n",
      "-15: 0.78\n",
      "=====\n",
      "-16: 0.766\n",
      "=====\n",
      "-17: 0.678\n",
      "=====\n",
      "-18: 0.638\n",
      "=====\n",
      "-19: 0.62\n",
      "=====\n",
      "-20: 0.616\n",
      "=====\n",
      "-21: 0.552\n",
      "=====\n",
      "-22: 0.444\n",
      "=====\n",
      "-23: 0.63\n",
      "=====\n",
      "-24: 0.474\n",
      "=====\n",
      "-25: 0.61\n",
      "=====\n",
      "-26: 0.588\n",
      "=====\n",
      "-27: 0.564\n",
      "=====\n",
      "-28: 0.44\n",
      "=====\n",
      "-29: 0.56\n",
      "=====\n",
      "-30: 0.622\n",
      "=====\n",
      "-31: 0.594\n",
      "=====\n",
      "-32: 0.41\n",
      "=====\n",
      "-33: 0.588\n",
      "=====\n",
      "-34: 0.416\n",
      "=====\n",
      "-35: 0.576\n",
      "=====\n",
      "-36: 0.49\n",
      "=====\n",
      "-37: 0.472\n",
      "=====\n",
      "-38: 0.486\n",
      "=====\n",
      "-39: 0.558\n",
      "=====\n",
      "-40: 0.566\n",
      "=====\n",
      "-41: 0.564\n",
      "=====\n",
      "-42: 0.564\n",
      "=====\n",
      "-43: 0.546\n",
      "=====\n",
      "-44: 0.49\n",
      "=====\n",
      "-45: 0.494\n",
      "=====\n",
      "-46: 0.484\n",
      "=====\n",
      "-47: 0.472\n",
      "=====\n"
     ]
    }
   ],
   "source": [
    "batch_size=128\n",
    "rep_reader = rep_pipeline.get_directions(\n",
    "    train_data,\n",
    "    rep_token=rep_token, \n",
    "    hidden_layers=hidden_layers, \n",
    "    n_difference=n_difference, \n",
    "    train_labels=train_labels, \n",
    "    direction_method=direction_method,\n",
    "    direction_finder_kwargs=direction_finder_kwargs,\n",
    "    batch_size=batch_size,\n",
    "    \n",
    "    max_length=2048,\n",
    "    padding=\"longest\",\n",
    ")\n",
    "\n",
    "batch_size=128\n",
    "results_val = {layer: {} for layer in hidden_layers}\n",
    "\n",
    "H_tests = rep_pipeline(\n",
    "                    test_data,\n",
    "                    rep_token=rep_token, \n",
    "                    hidden_layers=hidden_layers, \n",
    "                    rep_reader=rep_reader,\n",
    "                    batch_size=batch_size,\n",
    "                    max_length=2048,\n",
    "                    padding=\"longest\")\n",
    "\n",
    "n_choices=2\n",
    "for layer in hidden_layers:\n",
    "    H_test = [H[layer] for H in H_tests] \n",
    "    H_test = [H_test[i:i+n_choices] for i in range(0, len(H_test), n_choices)]\n",
    "\n",
    "    sign = rep_reader.direction_signs[layer]\n",
    "    eval_func = min if sign == -1 else max\n",
    "    \n",
    "    cors = np.mean([eval_func(H) == H[0] for H in H_test])\n",
    "    results_val[layer] = cors\n",
    "    \n",
    "    print(f\"{layer}: {cors}\")\n",
    "    print(\"=====\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cc6357ff-71a3-4e81-8542-956c2fc5a0a0",
   "metadata": {},
   "source": [
    "## QNLI"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "992d710e-ab5d-4118-a843-8ea939c33f6b",
   "metadata": {},
   "outputs": [],
   "source": [
    "template_str = 'Consider the {concept} of the answer to the question:\\nQuestion: {question}\\nAnswer: {sentence}'\n",
    "rep_token = 3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "d5028907-5c3e-4f88-90f9-27de5b5ccc1f",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Found cached dataset glue (/data/long_phan/.cache/huggingface/datasets/glue/qnli/1.0.0/dacbe3125aa31d7f70367a07a8a9e72a5a0bfeb5fc42e75c9db75b96da6053ad)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d49c74c563994d32b5201d784bed0e71",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached shuffled indices for dataset at /data/long_phan/.cache/huggingface/datasets/glue/qnli/1.0.0/dacbe3125aa31d7f70367a07a8a9e72a5a0bfeb5fc42e75c9db75b96da6053ad/cache-b7167382490e7670.arrow\n"
     ]
    }
   ],
   "source": [
    "def sample(ds, eval_set=False):\n",
    "    pairs = []\n",
    "    for e in ds:\n",
    "        concepts = ['plausibility', 'implausibility']\n",
    "        sentence  = e['sentence']\n",
    "        question = e['question']\n",
    "                \n",
    "        if e['label'] == 1 and eval_set: # flip the concepts in test for easy eval\n",
    "            concepts = concepts[::-1]\n",
    "        pair = [template_str.format(concept=c, question=question, sentence=sentence) for c in concepts]\n",
    "        pairs.extend(pair)\n",
    "    return pairs\n",
    "\n",
    "max_train_samples = 128\n",
    "\n",
    "dataset = load_dataset(\"glue\", \"qnli\")\n",
    "train_dataset, test_dataset = dataset['train'], dataset['validation']\n",
    "train_dataset =  train_dataset.shuffle(seed=0)\n",
    "train_dataset = train_dataset.select(range(max_train_samples))\n",
    "\n",
    "train_data, train_labels = sample(train_dataset), train_dataset['label']\n",
    "train_labels = [[1, 0] if label == 0 else [0, 1] for label in train_labels]\n",
    "\n",
    "test_data = sample(test_dataset, eval_set=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "b34d7521-38b0-49bb-9710-ba54ed1fcdd5",
   "metadata": {},
   "outputs": [],
   "source": [
    "rep_reader = rep_pipeline.get_directions(\n",
    "    train_data,\n",
    "    rep_token=rep_token, \n",
    "    hidden_layers=hidden_layers, \n",
    "    n_difference=n_difference, \n",
    "    train_labels=train_labels, \n",
    "    direction_method=direction_method,\n",
    "    direction_finder_kwargs=direction_finder_kwargs,\n",
    "    batch_size=8,\n",
    "    \n",
    "    max_length=2048,\n",
    "    padding=\"longest\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "83177364-abca-41bf-925c-a976c1bc3a3a",
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size=128\n",
    "results_val = {layer: {} for layer in hidden_layers}\n",
    "\n",
    "H_tests = rep_pipeline(test_data, \n",
    "                    rep_token=rep_token, \n",
    "                    hidden_layers=hidden_layers, \n",
    "                    rep_reader=rep_reader,\n",
    "                    batch_size=batch_size,\n",
    "                    max_length=2048,\n",
    "                    padding=\"longest\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "2b299e18-f356-4235-806b-f9551e2728a9",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-1: 0.6976020501555922\n",
      "=====\n",
      "-2: 0.4946000366099213\n",
      "=====\n",
      "-3: 0.5773384587223137\n",
      "=====\n",
      "-4: 0.6686802123375435\n",
      "=====\n",
      "-5: 0.6132161815852096\n",
      "=====\n",
      "-6: 0.5464030752333883\n",
      "=====\n",
      "-7: 0.5374336445176643\n",
      "=====\n",
      "-8: 0.5405454878272011\n",
      "=====\n",
      "-9: 0.5251693208859601\n",
      "=====\n",
      "-10: 0.5174812374153396\n",
      "=====\n",
      "-11: 0.5370675453047776\n",
      "=====\n",
      "-12: 0.5291964122277137\n",
      "=====\n",
      "-13: 0.5403624382207578\n",
      "=====\n",
      "-14: 0.5648910854841662\n",
      "=====\n",
      "-15: 0.5606809445359693\n",
      "=====\n",
      "-16: 0.556836902800659\n",
      "=====\n",
      "-17: 0.528281164195497\n",
      "=====\n",
      "-18: 0.5207761303313198\n",
      "=====\n",
      "-19: 0.49679663188724144\n",
      "=====\n",
      "-20: 0.48288486179754714\n",
      "=====\n",
      "-21: 0.47995606809445357\n",
      "=====\n",
      "-22: 0.4885593995972909\n",
      "=====\n",
      "-23: 0.48929159802306427\n",
      "=====\n",
      "-24: 0.4995423759838916\n",
      "=====\n",
      "-25: 0.5143693941058026\n",
      "=====\n",
      "-26: 0.5022881200805418\n",
      "=====\n",
      "-27: 0.5088779059125023\n",
      "=====\n",
      "-28: 0.5081457074867289\n",
      "=====\n",
      "-29: 0.4993593263774483\n",
      "=====\n",
      "-30: 0.5015559216547685\n",
      "=====\n",
      "-31: 0.5046677649643053\n",
      "=====\n",
      "-32: 0.48416620904265056\n",
      "=====\n",
      "-33: 0.4925864909390445\n",
      "=====\n",
      "-34: 0.4914881933003844\n",
      "=====\n",
      "-35: 0.5112575507962658\n",
      "=====\n",
      "-36: 0.5107084019769358\n",
      "=====\n",
      "-37: 0.5130880468606992\n",
      "=====\n",
      "-38: 0.5107084019769358\n",
      "=====\n",
      "-39: 0.5026542192934286\n",
      "=====\n",
      "-40: 0.5017389712612118\n",
      "=====\n",
      "-41: 0.5026542192934286\n",
      "=====\n",
      "-42: 0.49093904448105435\n",
      "=====\n",
      "-43: 0.5019220208676551\n",
      "=====\n",
      "-44: 0.5021050704740985\n",
      "=====\n",
      "-45: 0.48929159802306427\n",
      "=====\n",
      "-46: 0.5074135090609555\n",
      "=====\n",
      "-47: 0.5085118066996156\n",
      "=====\n"
     ]
    }
   ],
   "source": [
    "n_choices=2\n",
    "for layer in hidden_layers:\n",
    "    H_test = [H[layer] for H in H_tests] \n",
    "    H_test = [H_test[i:i+n_choices] for i in range(0, len(H_test), n_choices)]\n",
    "\n",
    "    sign = rep_reader.direction_signs[layer]\n",
    "    eval_func = min if sign == -1 else max\n",
    "    \n",
    "    cors = np.mean([eval_func(H) == H[0] for H in H_test])\n",
    "    results_val[layer] = cors\n",
    "    \n",
    "    print(f\"{layer}: {cors}\")\n",
    "    print(\"=====\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8cbcacf9-6ad4-4a05-9c6d-7534e0728c70",
   "metadata": {},
   "source": [
    "## PIQA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "5c9b198a-3b7b-49b9-add3-f3fbec02cf29",
   "metadata": {},
   "outputs": [],
   "source": [
    "template_str = \"Consider the amount of plausible reasoning in the scenario:\\n{goal} {sol}\"\n",
    "rep_token = 5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "4e3eadb9-8bc4-4577-a796-41e64f1fb962",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Found cached dataset piqa (/data/long_phan/.cache/huggingface/datasets/piqa/plain_text/1.1.0/6c611c1a9bf220943c4174e117d3b660859665baf1d43156230116185312d011)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b36e55f64cbd457e9753294a57f4c285",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached shuffled indices for dataset at /data/long_phan/.cache/huggingface/datasets/piqa/plain_text/1.1.0/6c611c1a9bf220943c4174e117d3b660859665baf1d43156230116185312d011/cache-ed09e117332c4dbb.arrow\n"
     ]
    }
   ],
   "source": [
    "def samples(ds, test_set=False):\n",
    "    pairs = []\n",
    "    for e in ds:\n",
    "        solutions  = [e['sol1'], e['sol2']]\n",
    "        goal = e['goal']\n",
    "                \n",
    "        if e['label'] == 1 and test_set: # flip the true sol in test for easy eval\n",
    "            solutions = solutions[::-1]\n",
    "\n",
    "        pair = [template_str.format(goal=goal, sol=sol) for sol in solutions]\n",
    "        pairs.extend(pair)\n",
    "    return pairs\n",
    "\n",
    "max_train_samples = 64\n",
    "\n",
    "dataset = load_dataset(\"piqa\")\n",
    "train_dataset, test_dataset = dataset['train'], dataset['validation']\n",
    "train_dataset =  train_dataset.shuffle(seed=0)\n",
    "train_dataset = train_dataset.select(range(max_train_samples))\n",
    "\n",
    "train_data, train_labels = samples(train_dataset), train_dataset['label']\n",
    "train_labels = [[1, 0] if label == 0 else [0, 1] for label in train_labels]\n",
    "\n",
    "test_data = samples(test_dataset, test_set=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "e9ae84f9-e751-499a-80a9-974b37f1f4be",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-1: 0.6996735582154516\n",
      "=====\n",
      "-2: 0.5712731229597389\n",
      "=====\n",
      "-3: 0.6964091403699674\n",
      "=====\n",
      "-4: 0.6991294885745375\n",
      "=====\n",
      "-5: 0.6936887921653971\n",
      "=====\n",
      "-6: 0.6866158868335147\n",
      "=====\n",
      "-7: 0.6877040261153428\n",
      "=====\n",
      "-8: 0.6942328618063112\n",
      "=====\n",
      "-9: 0.6893362350380848\n",
      "=====\n",
      "-10: 0.6877040261153428\n",
      "=====\n",
      "-11: 0.6871599564744287\n",
      "=====\n",
      "-12: 0.6964091403699674\n",
      "=====\n",
      "-13: 0.6947769314472253\n",
      "=====\n",
      "-14: 0.6708378672470077\n",
      "=====\n",
      "-15: 0.6942328618063112\n",
      "=====\n",
      "-16: 0.6822633297062024\n",
      "=====\n",
      "-17: 0.6343852013057671\n",
      "=====\n",
      "-18: 0.6066376496191512\n",
      "=====\n",
      "-19: 0.5990206746463548\n",
      "=====\n",
      "-20: 0.5794341675734495\n",
      "=====\n",
      "-21: 0.5413492927094669\n",
      "=====\n",
      "-22: 0.47551686615886835\n",
      "=====\n",
      "-23: 0.4591947769314472\n",
      "=====\n",
      "-24: 0.4733405875952122\n",
      "=====\n",
      "-25: 0.4820457018498368\n",
      "=====\n",
      "-26: 0.4619151251360174\n",
      "=====\n",
      "-27: 0.49347116430903154\n",
      "=====\n",
      "-28: 0.4923830250272035\n",
      "=====\n",
      "-29: 0.4836779107725789\n",
      "=====\n",
      "-30: 0.48639825897714906\n",
      "=====\n",
      "-31: 0.514145810663765\n",
      "=====\n",
      "-32: 0.5457018498367792\n",
      "=====\n",
      "-33: 0.45212187159956474\n",
      "=====\n",
      "-34: 0.4602829162132753\n",
      "=====\n",
      "-35: 0.5380848748639826\n",
      "=====\n",
      "-36: 0.47170837867247006\n",
      "=====\n",
      "-37: 0.4695321001088139\n",
      "=====\n",
      "-38: 0.5065288356909684\n",
      "=====\n",
      "-39: 0.499455930359086\n",
      "=====\n",
      "-40: 0.4967355821545158\n",
      "=====\n",
      "-41: 0.499455930359086\n",
      "=====\n",
      "-42: 0.5065288356909684\n",
      "=====\n",
      "-43: 0.47551686615886835\n",
      "=====\n",
      "-44: 0.5201305767138193\n",
      "=====\n",
      "-45: 0.47878128400435255\n",
      "=====\n",
      "-46: 0.4896626768226333\n",
      "=====\n",
      "-47: 0.49183895538628947\n",
      "=====\n"
     ]
    }
   ],
   "source": [
    "batch_size=128\n",
    "rep_reader = rep_pipeline.get_directions(\n",
    "    train_data,\n",
    "    rep_token=rep_token, \n",
    "    hidden_layers=hidden_layers, \n",
    "    n_difference=n_difference, \n",
    "    train_labels=train_labels, \n",
    "    direction_method=direction_method,\n",
    "    direction_finder_kwargs=direction_finder_kwargs,\n",
    "    batch_size=batch_size,\n",
    "    \n",
    "    max_length=2048,\n",
    "    padding=\"longest\",\n",
    ")\n",
    "\n",
    "batch_size=128\n",
    "results_val = {layer: {} for layer in hidden_layers}\n",
    "\n",
    "H_tests = rep_pipeline(\n",
    "                    test_data,\n",
    "                    rep_token=rep_token, \n",
    "                    hidden_layers=hidden_layers, \n",
    "                    rep_reader=rep_reader,\n",
    "                    batch_size=batch_size,\n",
    "                    max_length=2048,\n",
    "                    padding=\"longest\")\n",
    "\n",
    "n_choices=2\n",
    "for layer in hidden_layers:\n",
    "    H_test = [H[layer] for H in H_tests] \n",
    "    H_test = [H_test[i:i+n_choices] for i in range(0, len(H_test), n_choices)]\n",
    "\n",
    "    sign = rep_reader.direction_signs[layer]\n",
    "    eval_func = min if sign == -1 else max\n",
    "    \n",
    "    cors = np.mean([eval_func(H) == H[0] for H in H_test])\n",
    "    results_val[layer] = cors\n",
    "    \n",
    "    print(f\"{layer}: {cors}\")\n",
    "    print(\"=====\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "80bb6a4e-d43e-4a45-a087-46169bb35da8",
   "metadata": {},
   "source": [
    "## COPA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "d128002b-9a84-4391-907a-8da649196722",
   "metadata": {},
   "outputs": [],
   "source": [
    "template_str = \"Consider the amount of plausible reasoning in the scenario:\\n{premise} {hook} {alternative}\"\n",
    "rep_token = 5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "8429fded-d7c5-44b5-92a0-f660ab2901d6",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Found cached dataset csv (/data/long_phan/.cache/huggingface/datasets/pkavumba___csv/pkavumba--balanced-copa-81fb5dd3c6eeb4d7/0.0.0/eea64c71ca8b46dd3f537ed218fc9bf495d5707789152eb2764f5c78fa66d59d)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "8cad233a9f4b48579ef184d7f791b922",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/2 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at /data/long_phan/.cache/huggingface/datasets/pkavumba___csv/pkavumba--balanced-copa-81fb5dd3c6eeb4d7/0.0.0/eea64c71ca8b46dd3f537ed218fc9bf495d5707789152eb2764f5c78fa66d59d/cache-f6b83c6a27e35715.arrow\n",
      "Loading cached shuffled indices for dataset at /data/long_phan/.cache/huggingface/datasets/pkavumba___csv/pkavumba--balanced-copa-81fb5dd3c6eeb4d7/0.0.0/eea64c71ca8b46dd3f537ed218fc9bf495d5707789152eb2764f5c78fa66d59d/cache-b4e8b29247fbf014.arrow\n"
     ]
    }
   ],
   "source": [
    "copa_hook = {\n",
    "    'cause': 'because',\n",
    "    'effect': 'then',\n",
    "}\n",
    "def samples(ds, test_set=False):\n",
    "    pairs = []\n",
    "    for e in ds:\n",
    "        alternatives  = [e['choice1'], e['choice2']]\n",
    "        premise = e['premise']\n",
    "        hook = copa_hook[e['question']]\n",
    "                \n",
    "        if e['label'] == 1 and test_set: # flip the true alt in test for easy eval\n",
    "            alternatives = alternatives[::-1]\n",
    "        pair = [template_str.format(premise=premise, hook=hook, alternative=alternative) for alternative in alternatives]\n",
    "        pairs.extend(pair)\n",
    "    return pairs\n",
    "\n",
    "max_train_samples = 64\n",
    "dataset = load_dataset(\"pkavumba/balanced-copa\")\n",
    "# Dataset is from https://github.com/Balanced-COPA/Balanced-COPA\n",
    "# Need to filter out mirrored samples to get original COPA\n",
    "dataset['train'] =  dataset['train'].filter(lambda e: not e['mirrored'])\n",
    "\n",
    "train_dataset, test_dataset= dataset['train'], dataset['test']\n",
    "train_dataset =  train_dataset.shuffle(seed=0)\n",
    "train_dataset = train_dataset.select(range(max_train_samples))\n",
    "\n",
    "train_data, train_labels = samples(train_dataset), train_dataset['label']\n",
    "train_labels = [[1, 0] if label == 0 else [0, 1] for label in train_labels]\n",
    "\n",
    "test_data = samples(test_dataset, test_set=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "dca4db91-403f-42d4-83f2-0c731e002497",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-1: 0.906\n",
      "=====\n",
      "-2: 0.722\n",
      "=====\n",
      "-3: 0.896\n",
      "=====\n",
      "-4: 0.894\n",
      "=====\n",
      "-5: 0.892\n",
      "=====\n",
      "-6: 0.886\n",
      "=====\n",
      "-7: 0.868\n",
      "=====\n",
      "-8: 0.872\n",
      "=====\n",
      "-9: 0.866\n",
      "=====\n",
      "-10: 0.872\n",
      "=====\n",
      "-11: 0.87\n",
      "=====\n",
      "-12: 0.862\n",
      "=====\n",
      "-13: 0.868\n",
      "=====\n",
      "-14: 0.856\n",
      "=====\n",
      "-15: 0.86\n",
      "=====\n",
      "-16: 0.85\n",
      "=====\n",
      "-17: 0.786\n",
      "=====\n",
      "-18: 0.746\n",
      "=====\n",
      "-19: 0.748\n",
      "=====\n",
      "-20: 0.732\n",
      "=====\n",
      "-21: 0.732\n",
      "=====\n",
      "-22: 0.724\n",
      "=====\n",
      "-23: 0.736\n",
      "=====\n",
      "-24: 0.668\n",
      "=====\n",
      "-25: 0.682\n",
      "=====\n",
      "-26: 0.702\n",
      "=====\n",
      "-27: 0.702\n",
      "=====\n",
      "-28: 0.682\n",
      "=====\n",
      "-29: 0.644\n",
      "=====\n",
      "-30: 0.614\n",
      "=====\n",
      "-31: 0.562\n",
      "=====\n",
      "-32: 0.544\n",
      "=====\n",
      "-33: 0.546\n",
      "=====\n",
      "-34: 0.532\n",
      "=====\n",
      "-35: 0.514\n",
      "=====\n",
      "-36: 0.512\n",
      "=====\n",
      "-37: 0.518\n",
      "=====\n",
      "-38: 0.544\n",
      "=====\n",
      "-39: 0.59\n",
      "=====\n",
      "-40: 0.568\n",
      "=====\n",
      "-41: 0.57\n",
      "=====\n",
      "-42: 0.564\n",
      "=====\n",
      "-43: 0.538\n",
      "=====\n",
      "-44: 0.54\n",
      "=====\n",
      "-45: 0.546\n",
      "=====\n",
      "-46: 0.566\n",
      "=====\n",
      "-47: 0.542\n",
      "=====\n"
     ]
    }
   ],
   "source": [
    "batch_size=128\n",
    "rep_reader = rep_pipeline.get_directions(\n",
    "    train_data,\n",
    "    rep_token=rep_token, \n",
    "    hidden_layers=hidden_layers, \n",
    "    n_difference=n_difference, \n",
    "    train_labels=train_labels, \n",
    "    direction_method=direction_method,\n",
    "    direction_finder_kwargs=direction_finder_kwargs,\n",
    "    batch_size=batch_size,\n",
    "    \n",
    "    max_length=2048,\n",
    "    padding=\"longest\",\n",
    ")\n",
    "\n",
    "batch_size=128\n",
    "results_val = {layer: {} for layer in hidden_layers}\n",
    "\n",
    "H_tests = rep_pipeline(\n",
    "                    test_data,\n",
    "                    rep_token=rep_token, \n",
    "                    hidden_layers=hidden_layers, \n",
    "                    rep_reader=rep_reader,\n",
    "                    batch_size=batch_size,\n",
    "                    max_length=2048,\n",
    "                    padding=\"longest\")\n",
    "\n",
    "n_choices=2\n",
    "for layer in hidden_layers:\n",
    "    H_test = [H[layer] for H in H_tests] \n",
    "    H_test = [H_test[i:i+n_choices] for i in range(0, len(H_test), n_choices)]\n",
    "\n",
    "    sign = rep_reader.direction_signs[layer]\n",
    "    eval_func = min if sign == -1 else max\n",
    "    \n",
    "    cors = np.mean([eval_func(H) == H[0] for H in H_test])\n",
    "    results_val[layer] = cors\n",
    "    \n",
    "    print(f\"{layer}: {cors}\")\n",
    "    print(\"=====\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "36c341d2-dbd2-4aa7-a2f9-761992eaf22e",
   "metadata": {},
   "source": [
    "## Story Cloze"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "c9cbaaaa-f4f7-4dcc-a6ce-4b6e46446685",
   "metadata": {},
   "outputs": [],
   "source": [
    "template_str = \"Consider the plausibility of the scenario:\\n{scenario}\"\n",
    "rep_token = 3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "756eca03-29a8-4f7d-bea9-e5948b5e76be",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Found cached dataset story_cloze (/data/long_phan/.cache/huggingface/datasets/story_cloze/2016-4349b0206f129d71/0.0.0/45cead0538c3deb72d731a7990e60835c2c9c5d5d5b1e95a7dd47ccf593671e4)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4f9573bf57ee4ac0a1853e202feba166",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/2 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached shuffled indices for dataset at /data/long_phan/.cache/huggingface/datasets/story_cloze/2016-4349b0206f129d71/0.0.0/45cead0538c3deb72d731a7990e60835c2c9c5d5d5b1e95a7dd47ccf593671e4/cache-cd5f729fe91c7970.arrow\n"
     ]
    }
   ],
   "source": [
    "def samples(ds, test_set=False):\n",
    "    pairs = []\n",
    "    for e in ds:\n",
    "        concepts = ['plausibility']\n",
    "        sentences = ' '.join([e[f'input_sentence_{i}'] for i in range(1, 5)])\n",
    "        quiz_sentences = [e['sentence_quiz1'], e['sentence_quiz2']]\n",
    "                \n",
    "        if e['answer_right_ending'] == 2 and test_set: # flip to have true sentence on first in test set for easy evaluation \n",
    "            quiz_sentences = quiz_sentences[::-1]\n",
    "\n",
    "        pair = [template_str.format(scenario=' '.join([sentences, s])) for s in quiz_sentences]\n",
    "        pairs.extend(pair)\n",
    "    return pairs\n",
    "\n",
    "max_train_samples = 64\n",
    "\n",
    "data_dir=\"\" # See https://huggingface.co/datasets/story_cloze to download data manually\n",
    "data_dir='/data/private_models/cais_models/misc/'\n",
    "dataset = load_dataset(\"story_cloze\", '2016', data_dir=data_dir)\n",
    "\n",
    "train_dataset, test_dataset = dataset['validation'], dataset['test']\n",
    "train_dataset =  train_dataset.shuffle(seed=0)\n",
    "train_dataset = train_dataset.select(range(max_train_samples))\n",
    "\n",
    "train_data = samples(train_dataset)\n",
    "train_labels = [[1, 0] if label == 1 else [0, 1] for label in train_dataset['answer_right_ending']]\n",
    "\n",
    "test_data = samples(test_dataset, test_set=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "6ae99e03-62f7-4327-a397-3d371c89e37c",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-1: 0.9716729021913415\n",
      "=====\n",
      "-2: 0.9294494922501336\n",
      "=====\n",
      "-3: 0.9652592196686264\n",
      "=====\n",
      "-4: 0.9700694815606627\n",
      "=====\n",
      "-5: 0.9679315873864244\n",
      "=====\n",
      "-6: 0.965793693212186\n",
      "=====\n",
      "-7: 0.9529663281667557\n",
      "=====\n",
      "-8: 0.9497594869053981\n",
      "=====\n",
      "-9: 0.9492250133618386\n",
      "=====\n",
      "-10: 0.9428113308391235\n",
      "=====\n",
      "-11: 0.943345804382683\n",
      "=====\n",
      "-12: 0.9363976483164084\n",
      "=====\n",
      "-13: 0.9390700160342063\n",
      "=====\n",
      "-14: 0.928915018706574\n",
      "=====\n",
      "-15: 0.9208979155531801\n",
      "=====\n",
      "-16: 0.8727952966328166\n",
      "=====\n",
      "-17: 0.8551576696953501\n",
      "=====\n",
      "-18: 0.8268305718866916\n",
      "=====\n",
      "-19: 0.8001068947087119\n",
      "=====\n",
      "-20: 0.8295029396044896\n",
      "=====\n",
      "-21: 0.8134687332977018\n",
      "=====\n",
      "-22: 0.8348476750400855\n",
      "=====\n",
      "-23: 0.8107963655799038\n",
      "=====\n",
      "-24: 0.7771245323356494\n",
      "=====\n",
      "-25: 0.7648316408337787\n",
      "=====\n",
      "-26: 0.772314270443613\n",
      "=====\n",
      "-27: 0.7509353287012293\n",
      "=====\n",
      "-28: 0.7482629609834314\n",
      "=====\n",
      "-29: 0.7274184927846071\n",
      "=====\n",
      "-30: 0.7541421699625869\n",
      "=====\n",
      "-31: 0.7509353287012293\n",
      "=====\n",
      "-32: 0.7338321753073223\n",
      "=====\n",
      "-33: 0.7274184927846071\n",
      "=====\n",
      "-34: 0.6932121859967931\n",
      "=====\n",
      "-35: 0.6900053447354356\n",
      "=====\n",
      "-36: 0.692143238909674\n",
      "=====\n",
      "-37: 0.6958845537145911\n",
      "=====\n",
      "-38: 0.6942811330839124\n",
      "=====\n",
      "-39: 0.6900053447354356\n",
      "=====\n",
      "-40: 0.6830571886691609\n",
      "=====\n",
      "-41: 0.6702298236237306\n",
      "=====\n",
      "-42: 0.6659540352752539\n",
      "=====\n",
      "-43: 0.632816675574559\n",
      "=====\n",
      "-44: 0.6322822020309995\n",
      "=====\n",
      "-45: 0.6365579903794762\n",
      "=====\n",
      "-46: 0.6221272047033671\n",
      "=====\n",
      "-47: 0.6109032602886157\n",
      "=====\n"
     ]
    }
   ],
   "source": [
    "batch_size=128\n",
    "rep_reader = rep_pipeline.get_directions(\n",
    "    train_data,\n",
    "    rep_token=rep_token, \n",
    "    hidden_layers=hidden_layers, \n",
    "    n_difference=n_difference, \n",
    "    train_labels=train_labels, \n",
    "    direction_method=direction_method,\n",
    "    direction_finder_kwargs=direction_finder_kwargs,\n",
    "    batch_size=batch_size,\n",
    "    \n",
    "    max_length=2048,\n",
    "    padding=\"longest\",\n",
    ")\n",
    "\n",
    "batch_size=128\n",
    "results_val = {layer: {} for layer in hidden_layers}\n",
    "\n",
    "H_tests = rep_pipeline(\n",
    "                    test_data,\n",
    "                    rep_token=rep_token, \n",
    "                    hidden_layers=hidden_layers, \n",
    "                    rep_reader=rep_reader,\n",
    "                    batch_size=batch_size,\n",
    "                    max_length=2048,\n",
    "                    padding=\"longest\")\n",
    "\n",
    "n_choices=2\n",
    "for layer in hidden_layers:\n",
    "    H_test = [H[layer] for H in H_tests] \n",
    "    H_test = [H_test[i:i+n_choices] for i in range(0, len(H_test), n_choices)]\n",
    "\n",
    "    sign = rep_reader.direction_signs[layer]\n",
    "    eval_func = min if sign == -1 else max\n",
    "    \n",
    "    cors = np.mean([eval_func(H) == H[0] for H in H_test])\n",
    "    results_val[layer] = cors\n",
    "    \n",
    "    print(f\"{layer}: {cors}\")\n",
    "    print(\"=====\")\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
